package com.cxx.purchasecharge.core.crypto;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.security.Key;
import java.security.MessageDigest;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.apache.log4j.Logger;

import com.cxx.purchasecharge.core.util.HexConvert;

/**
 * This class provides support for encrypting and decrypting data using AES-128.
 * The algorithm used ensures that encrypting the same plain text at a later
 * time, will produce different cipher text.
 * <p>
 * Note: In cryptography, the Advanced Encryption Standard (AES) is an
 * encryption standard adopted by the U.S. government. The standard comprises
 * three block ciphers, AES-128, AES-192 and AES-256. This class use AES-128.
 * </p>
 * <p>
 * Callers should use the CryptoFactory getInstance methods to obtain an implementation.
 * </p> 
 *
 * 
 * This code was copied from a UMA project by Marc Lottman.
 * 
 */
public class AESEncryption implements ICrypto
{
    private static final Logger logger = Logger.getLogger(AESEncryption.class);
    
    private static final String ALGORITHM = "AES";
    private static final String CIPHER_TYPE = "AES/CBC/PKCS5Padding";
    private static final String DIGEST_TYPE = "SHA-256";

    private static final byte[] DEFAULT_SALT = { 60, 0, 35, 71, 1, 60, 32, 91, 122, 121, 67, 84, 79, 93, 52, 75 };
    private static final byte[] DEFAULT_IV_SUFFIX = { -7, -48, -54, -118, -5, 98, -11, -45 };
    private static final String UTF8 = "UTF-8";
     
    private static final byte[] s_keyData = { (byte)0xF9, (byte)98, (byte)104, (byte)79, (byte)93, (byte)63, 
                                                 (byte)0xE1, (byte)1, (byte)117, (byte)10, (byte)37, (byte)0xC4, 
                                                 (byte)25, (byte)7, (byte)0xF5, (byte)0xEE };
    
    protected Key m_key;

    /**
     * Creates an instance of the AESEncryption class using the default key text and
     * default salt.
     * 
     * Callers should use the CryptoFactory getInstance methods to obtain an instance
     * of this implementation.
     * 
     * @throws CryptoException
     *             if there are any problems obtaining the required encryption
     *             algorithms.
     */
    public AESEncryption() throws CryptoException
    {
        try
        {
            m_key = new SecretKeySpec (s_keyData, ALGORITHM);

            // Check that we can get the appropriate Cipher and digest.
            // We don't actually use or save them here, but want to error
            // quickly if for some strange reason the ciphers are not available.
            Cipher.getInstance (CIPHER_TYPE);
            MessageDigest.getInstance (DIGEST_TYPE);
            
        }
        catch (Exception e)
        {
            throw new CryptoException (e);
        }
    }
    
    public Key getKey()
    {
        return m_key;
    }
    
    /**
     * Creates an instance of the AESEncryption class using the specified key text and
     * the default salt.
     * 
     * Callers should use the CryptoFactory getInstance methods to obtain an instance
     * of this implementation.
     * 
     * @param keyText The key used to encrypt the data.
     * @throws CryptoException
     *             if there are any problems obtaining the required encryption
     *             algorithms.
     * @throws IllegalArgumentException
     *             if keyText is null
     */
    public AESEncryption(String keyText) throws CryptoException
    {
        this(keyText, DEFAULT_SALT);
    }
    
    /**
     * Creates an instance of the AESEncryption class using the specified key and salt.
     * 
     * Callers should use the CryptoFactory getInstance methods to obtain an instance
     * of this implementation.
     * 
     * @param keyText
     *            The key used to encrypt the text.
     * @param keySalt
     *            The salt used by the encryption/decryption.
     * @throws CryptographyException
     *             if there are any problems obtaining the required encryption
     *             algorithms.
     * @throws IllegalArgumentException
     *             if keyText or keySalt is null 
     */
    public AESEncryption(String keyText, byte[] keySalt) throws CryptoException
    {
        
        if ( keyText == null )
        {
            throw new IllegalArgumentException("keyText cannot be null");
        }
        
        if ( keySalt == null )
        {
            throw new IllegalArgumentException("keySalt cannot be null");
        }
        
        try
        {            
            // Derive a key from the key text and salt.
            // The idea here is to make it more computationally expensive to
            // test various keyTexts, since for each guess you would have to run it
            // thru the below algorithm.  In addition it allows us to create
            // the size of the key we need.  Java provides a PBKDF2WithHmacSHA1
            // function, to derive the key, but unfortunately it is not available in 
            // IBM Java 5.
            MessageDigest md = MessageDigest.getInstance (DIGEST_TYPE);           
            byte[] keyData = keyText.getBytes (UTF8);   
            
            // Recursively hash - adding the hash result plus the salt
            // together each time.
            for (int i = 0; i < 16000; i++)
            {
                // normally one would create the data byte array outside
                // of the loop, but the size of it varies in the first
                // call since the initial keyData is just the bytes 
                // from the keyText.  The rest of the calls have a
                // keyData that is the size of the digest, which is static.
                byte[] data = new byte[keyData.length + keySalt.length];             
                System.arraycopy (keyData, 0, data, 0, keyData.length);
                System.arraycopy (keySalt, 0, data, keyData.length, keySalt.length);
                keyData = md.digest (data);        
            }
            
            // Create a secret key with the first 16 bytes (128 bits) 
            // of the keyData.
            m_key = new SecretKeySpec (keyData, 0, 16, ALGORITHM);

            // Check that we can get the appropriate Cipher.
            // We don't actually use the Cipher here, but want to error
            // quickly if for some strange reason the cipher is not available.
            Cipher.getInstance (CIPHER_TYPE);  
            
        }
        catch (Exception e)
        {
            throw new CryptoException (e);
        }
    }
    
    /**
     * Generate initialization vector for AES encryption.
     * 
     * @return Initialization Vector.
     */
    private static byte[] createIV ()
    {
        ByteArrayOutputStream s;
        try
        {
            s = new ByteArrayOutputStream ();
            DataOutputStream out = new DataOutputStream (s);
            out.writeLong (System.currentTimeMillis ());
            out.write (DEFAULT_IV_SUFFIX);
            out.flush ();
            out.close ();
            return s.toByteArray ();
        }
        catch (IOException e)
        {
            // exceptions are not raised on byte array output
            // streams, so this should never happen.
            // Just in case we convert it to a runtime exception
            throw new RuntimeException (e);
        }
    }    
    
    /* (non-Javadoc)
     * @see net.carefx.crypto.ICrypto#decrypt(java.lang.String)
     */
    public String decrypt (String encryptedData) throws CryptoException
    {
        if ( encryptedData == null )
        {
            throw new IllegalArgumentException("encryptedData cannot be null");
        }
        
        try
        {
            byte[] data = HexConvert.convertToBytes(encryptedData);
            byte[] decryptedData = decrypt(data);
           
            return new String(decryptedData);
        }
        catch (CryptoException e)
        {
            throw e;
        }
        catch (Exception e)
        {
            throw new CryptoException (e);
        }
    }
    
    /* (non-Javadoc)
     * @see net.carefx.crypto.ICrypto#decrypt(byte[])
     */
    public byte[] decrypt (byte[] encryptedData) throws CryptoException
    {
        if ( encryptedData == null )
        {
            throw new IllegalArgumentException("encryptedData cannot be null");
        }
        
        try
        {
            // pull off the initialization vector, it is on the front of the
            // data.
            
            if (encryptedData.length < 16)
            {
                throw new CryptoException ("Encrypted data does not contain necessary initialization vector");
            }
            byte[] iv = new byte[16];
            System.arraycopy (encryptedData, 0, iv, 0, 16);

            Cipher cipher = Cipher.getInstance (CIPHER_TYPE);
            IvParameterSpec ivSpec = new IvParameterSpec(iv);
            cipher.init(Cipher.DECRYPT_MODE,getKey(),ivSpec);
            
            byte[] dataWithHash = cipher.doFinal (encryptedData, 16, encryptedData.length - 16);
            if (dataWithHash.length < 32)
            {
                throw new CryptoException ("Encrypted data does not contain necessary hash");
            }

            byte[] savedDigest = new byte[32];
            System.arraycopy (dataWithHash, 0, savedDigest, 0, 32);

            byte[] decryptedData = new byte[dataWithHash.length - 32];
            System.arraycopy (dataWithHash, 32, decryptedData, 0, decryptedData.length);

            // Create our own digest to compare
            MessageDigest md = MessageDigest.getInstance (DIGEST_TYPE);
            byte[] digest = md.digest (decryptedData);

            if (!MessageDigest.isEqual (savedDigest, digest))
            {
                throw new CryptoException ("Encrypted data has been tampered with.  Hashes do not match.");
            }

            return decryptedData;
        }
        catch (CryptoException e)
        {
            throw e;
        }
        catch (Exception e)
        {
            throw new CryptoException (e);
        }
    }
    /* (non-Javadoc)
     * @see net.carefx.crypto.ICrypto#decryptObject(byte[])
     */
    public Object decryptObject (byte[] data) throws CryptoException
    {
        if ( data == null )
        {
            throw new IllegalArgumentException("data cannot be null");
        }
        
        ObjectInputStream in = null;
        try
        {
            byte[] decrypted = decrypt (data);
            in = new ObjectInputStream (new ByteArrayInputStream (decrypted));
            
            return in.readObject ();
        }
        catch (Exception e)
        {
            throw new CryptoException (e);
        }
        finally
        {
            if (in != null)
            {
                try
                {
                    in.close ();
                }
                catch (IOException ignored)
                {
                }
            }
        }
    }
    
    /* (non-Javadoc)
     * @see net.carefx.crypto.ICrypto#encrypt(java.lang.String)
     */
    public String encrypt(String data) throws CryptoException
    {
        if ( data == null )
        {
            throw new IllegalArgumentException("data cannot be null");
        }
        
        try
        {
            // We create a hash of the data 32 bytes, and pre-pend it to the
            // data to encrypt. This allows us to verify in the
            // decrypt that the data has not be changed in any way.
            byte[] bData = data.getBytes(UTF8);
            byte[] results = encrypt(bData);   

            return HexConvert.convertHexString(results);
        }
        catch (Exception e)
        {
            throw new CryptoException (e);
        }
    }
    
    /* (non-Javadoc)
     * @see net.carefx.crypto.ICrypto#encrypt(byte[])
     */
    public byte[] encrypt(byte[] data) throws CryptoException
    {
        if ( data == null )
        {
            throw new IllegalArgumentException("data cannot be null");
        }
        
        try
        {
            // We create a hash of the data 32 bytes, and pre-pend it to the
            // data to encrypt. This allows us to verify in the
            // decrypt that the data has not be changed in any way.
            MessageDigest md = MessageDigest.getInstance (DIGEST_TYPE);
            byte[] digest = md.digest (data);
            byte[] iv = createIV();            
            byte[] dataWithHash = new byte[digest.length + data.length];
            System.arraycopy (digest, 0, dataWithHash, 0, digest.length);
            System.arraycopy (data, 0, dataWithHash, digest.length, data.length);

            Cipher cipher = Cipher.getInstance (CIPHER_TYPE);
            cipher.init (Cipher.ENCRYPT_MODE, getKey(), new IvParameterSpec (iv));

            byte[] encrypted = cipher.doFinal (dataWithHash);
            byte[] results = new byte[encrypted.length + 16];
            System.arraycopy (iv, 0, results, 0, 16);
            System.arraycopy (encrypted, 0, results, 16, encrypted.length);
            
            return results;
        }
        catch (Exception e)
        {
            throw new CryptoException (e);
        }
    }
    
    /* (non-Javadoc)
     * @see net.carefx.crypto.ICrypto#encryptObject(java.lang.Object)
     */
    public byte[] encryptObject (Object obj) throws CryptoException
    {
        if ( obj == null )
        {
            throw new IllegalArgumentException("obj cannot be null");
        }
        
        ObjectOutputStream out = null;
        try
        {
            ByteArrayOutputStream bytes = new ByteArrayOutputStream ();
            out = new ObjectOutputStream (bytes);
            out.writeObject (obj);
            out.close ();
            
            byte[] retb = encrypt(bytes.toByteArray());
            
            return retb;
        }
        catch (Exception e)
        {
            throw new CryptoException (e);
        }
        finally
        {
            if (out != null)
            {
                try
                {
                    out.close ();
                }
                catch (IOException ignored)
                {
                }
            }
        }
    }
}
